import argparse
import json
from json.decoder import JSONDecodeError
import time
import os
import flask
from flask import request
from FlagEmbedding import BGEM3FlagModel
from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM
import requests
import logging
app = flask.Flask(__name__)
embedding = None
tokenizer_map = {
    # "model_name" :  tokenizer
}
from args import *
def construct_template(model_name, instruct):
    messages = [
            {"role": "user", "content": f"{instruct}"}
        ]
    if "chatglm3" in model_name:
        messages = [
            {"role": "<|system|>", "content": "你是一个有用的助手。"},
            {"role": "<|user|>", "content": f"{instruct}"}
            # {"role": "<|assistant|>", "content": f"}
        ]
    elif "qwen1.5" in model_name:
        messages = [
            {"role": "system", "content": "你是一个有用的助手。"},
            {"role": "user", "content": f"{instruct}"}
            # {"role": "assistant", "content": f"}
        ]
    elif "baichuan2" in model_name:
        messages = [
            {"role": "<reserved_106>", "content": f"{instruct}"}
            # {"role": "<reserved_107>", "content": f"}
        ]
    elif "internlm2" in model_name:
        messages = [
            {"role": "system", "content": "你是一个有用的助手。"},
            {"role": "user", "content": f"{instruct}"}
            # {"role": "assistant", "content": f"}
        ]
    elif "deepseek" in model_name:
        messages = [
            {"role": "User", "content": f"{instruct}"}
            # {"role": "Assistant", "content": f"}
        ]
    elif "yi" in model_name:
        messages = [
            {"role": "user", "content": f"{instruct}"}
            # {"role": "assistant", "content": f"}
        ]
    elif "mistral" in model_name:
        messages = [
            {"role": "[INST]", "content": "你是一个有用的助手。"},
            {"role": "[INST]", "content": f"{instruct}"}
            # {"role": "[/INST]", "content": f"}
        ]
    elif "llama2" in model_name:
        messages = [
            {"role": "[INST] <<SYS>>", "content": "你是一个有用的助手。"},
            {"role": "[INST]", "content": f"{instruct}"}
            # {"role": "[/INST]", "content": f"}
        ]
    elif "llama3" in model_name:
        messages = [
            {"role": "system", "content": "你是一个有用的助手。"},
            {"role": "user", "content": f"{instruct}"}
            # {"role": "assistant", "content": f"}
        ]
    elif "gemma" in model_name:
        messages = [
            {"role": "user", "content": f"{instruct}"}
        ]
    return messages

def do_request(model_name, messages, max_new_tokens=512):
    tokenizer = tokenizer_map.get(model_name, None)
    if tokenizer is None:
        tokenizer = AutoTokenizer.from_pretrained(f"{MODEL_PATH}/{model_name}", trust_remote_code=True, use_fast=False)
        tokenizer_map[model_name] = tokenizer
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    resp = requests.post("http://127.0.0.1:8898/control/call", json={"model": model_name, "text": text, "max_new_tokens": max_new_tokens})
    data = json.loads(resp.text)
    result = data.get("result", "")
    # result = result[len(text):]
    return result

def get_area(model_name, news, ):
    instruct = f"【领域】[{','.join(area)}]\n【文章】{news}\n【分类】文章中可能会有无关、垃圾内容，请忽略，请按照领域输出分类结果："
    messages = construct_template(model_name, instruct)# 微调模板 150条 1000条
    with open("log.txt", "a+") as f:
        f.write(instruct + "\n")
    text = do_request(model_name, messages)
    return text



def update_record(post_data, result):
    fname = os.path.join(
        "../record", time.strftime("data_%Y_%m_%d.json", time.localtime()))
    data = []
    if os.path.exists(fname):
        with open(fname, "r+") as f:
            data = json.load(f)
    with open(fname, "w+") as f:
        data.append({"post_data": post_data, "summary": result})
        json.dump(data, f, ensure_ascii=False, indent=4)
    return


def filter(newslist):
    newslist = [news for news in newslist if len(news) >= 128] 
    if len(newslist) == 0:
        return []
    # 100 news -> 100*1024 vect
    # vect.T * vect  100*100 
    vects = embedding.encode(newslist, batch_size=12, max_length=8192, )['dense_vecs']
    keylist, keyset = [], set()
    for i, v in enumerate(vects):
        if sum([(vects[k]@v) >= 0.82 for k in keyset]) == 0:
            keylist.append(newslist[i])
            keyset.add(i)
    return keylist


@app.route('/<string:model_name>/chat', methods=['POST'])
def show_post(model_name):
    try:
        abstract = request.json.get("news")
        result = get_area(model_name, abstract )
        # update_record(request.json, result)
        return json.dumps({"code": 0, "message": "操作成功！", "data": result}, indent=4, ensure_ascii=False)
    except Exception as e:
        return {"code": -1, "message": f"{e}", "data": ""}

#  http://10.26.32.91/record/data_2022_03_12.json


@app.route('/record/<string:file_name>', methods=['GET'])
def show_file(file_name):
    try:
        with open(os.path.join("../record", file_name)) as f:
            j = json.load(f)
        return {"code": 0, "message": "操作成功！", "data": j}
    except Exception as e:
        return {"code": -2, "message": f"{e}", "data": ""}


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=8900)
    args = parser.parse_args()

    embedding = BGEM3FlagModel(f'{MODEL_PATH}/bge-m3', use_fp16=True, device="cuda:6")

    app.run(host="0.0.0.0", port=args.port)
